from flask import Flask, render_template, Response, request, url_for,stream_with_context, jsonify
from skimage.io import imread, imsave
import os
import wave
from pathlib import Path
import numpy as np

from DataHidingDemo.extra import file_to_bin
from DataHidingDemo.extra import file_to_bin, mes_to_bin

app = Flask(__name__)
import time
# from werkzeug.utils import secure_filename
import random
import json
from DataHidingDemo.LSB_hide import lsb_hide, wav_lsb_hide, psnr, falseRate
from DataHidingDemo.LSB_extract import lsb_reveal, wav_lsb_extract
from DataHidingDemo.pvd import PVD
from DataHidingDemo.pvd import progress_bar_ratio
from DataHidingDemo.DCT_hide import dct_hide
# from werkzeug.datastructures import FileStorage
from DataHidingDemo.DCT_extract import dct_reveal
from DataHidingDemo.rs import stegRate
from DataHidingDemo.kf import stgPrb
from PIL import Image
import numpy as np
from matplotlib import image
from skimage import io, color
import logging
import timeit


# TODO: Flask 模块测试 Demo: https://flask.net.cn/testing.html

ALLOWED_EXTENSIONS = set(['png', 'PNG', 'bmp', 'wav'])

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS


def get_secrets_and_coverimg(request):
    if request.files:
        f = request.files['file1']
        logging.info(f"接收文件: {f.filename}")
        if not (f and allowed_file(f.filename)):
            logging.error("文件格式错误")
            print("file error")
            # return render_template('lsb.html', message="上传文件格式错误")
            return False, {'flag':'1', 'message': "文件格式错误"}
    else:
        logging.warning("没有文件上传")
        # return render_template('lsb.html', message="文件为空")
        return False, {'flag':'2', 'message':"没有文件上传"}
    secrets_text = request.form.get('secret_text')
    secrets_file = request.files.get('secret_file')
    if not secrets_text:
        secrets_text = secrets_file.read().decode('utf-8')
    return True, [f, secrets_text]

def prepareFile(file):
    result_file_pid = file.filename
    basepath = os.path.dirname(__file__)
    
    rf_name = os.path.join('data/', time.strftime("%Y.%m.%d"))
    result_folder_name = os.path.join(basepath, 'static', rf_name)
    # 对于每一个时间单独建立一个文件夹
    if os.path.exists(result_folder_name) and os.path.isdir(result_folder_name):
        existing_folders = [f for f in os.listdir(os.path.join(basepath, result_folder_name)) if
                            os.path.isdir(os.path.join(basepath, result_folder_name, f))]
        folder_index = len(existing_folders) + 1
    # 遍历存在过的文件夹避免重复
    else:
        existing_folders = 0
        folder_index = 1
        os.makedirs(result_folder_name)

    fold_name = str(folder_index)
    while (os.path.exists(fold_name) and os.path.isdir(fold_name)):
        # 以防出现同名错误
        folder_index = folder_index + 1
        fold_name = str(folder_index)
    tarPath = os.path.join(result_folder_name, fold_name)
    os.makedirs(tarPath)
    
    fileSuffix = result_file_pid.split('.')[-1]
    upload_path = os.path.join(result_folder_name, fold_name, "original."+fileSuffix)
    file.save(upload_path)
    
    config = {"id": folder_index, "dataDir":os.path.join(rf_name, fold_name), "originalPath": os.path.join(rf_name, fold_name, "original."+fileSuffix)}
    
    if fileSuffix == 'wav':
        with wave.open(upload_path, 'rb') as wav_in:
            # 获取WAV文件的参数
            params = wav_in.getparams()

            # 读取WAV文件中的音频数据
            frames = wav_in.readframes(wav_in.getnframes())
            # 保存音频数据到新的WAV文件
            with wave.open(os.path.join(basepath, 'static/wav', 'test.wav'), 'wb') as wav_out:
                wav_out.setparams(params)
                wav_out.writeframes(frames)
        config.update({"fileType": "wav", "fileSuffix": fileSuffix, "params":params})
        return frames, config
    else:
        img = imread(upload_path)
        # imsave(os.path.join(basepath, 'static/images', 'test.png'), img)
        config.update({"fileType": "img", "fileSuffix": fileSuffix})
        return img, config

def afterHide(stegData, config, img = None):
    basepath = os.path.dirname(__file__)
    if config.get("fileType") == "wav":
        dealed_wav_pid = str(int(time.time())) + str(random.randint(0, 10000)).zfill(5) + '_lsb.' + config.get("fileSuffix")
        dealed_path = os.path.join(os.path.dirname(config.get("originalPath")), dealed_wav_pid)
        stego_wav = wave.open(os.path.join(basepath, "static", dealed_path), 'wb')
        stego_wav.setparams(config.get("params"))
        stego_wav.writeframes(stegData)
        config["stegPath"] = dealed_path
        with open(os.path.join(basepath, 'static', config["dataDir"], 'config.json'), 'w') as f:
            json.dump(config, f)
        return config
    else:
        dealed_image_pid = config.get("originalFilename", "tmp") + str(int(time.time())) + str(random.randint(0, 10000)).zfill(5) + '_' + config.get("algName") + "." + config.get("fileSuffix")
        # 加上前缀以便识别宽高
        fold_name = str(config.get("id"))
        dealed_path = os.path.join(os.path.dirname(config.get("originalPath")), dealed_image_pid)
        
        imsave(os.path.join(basepath, "static", dealed_path), stegData)
        config["stegPath"] = dealed_path
        
        str_info = "转换成功"

        beforeLong=img.shape[0]
        beforeWide=img.shape[1]
        beforeChannel=3 if len(img.shape) == 3 else 1

        afterLong=stegData.shape[0]
        afterWide=stegData.shape[1]
        afterChannel=3 if len(stegData.shape) == 3 else 1

        condition = img!=stegData
        ids=np.where(condition)
        ids=np.column_stack(ids).tolist()

        if len(stegData.shape)==3:
            psnr_ = psnr(stegData, img)
            falseRate_ = falseRate(stegData, img)
        else:
            #TODO: 计算灰度图像的psnr和falseRate
            psnr_ = 0
            falseRate_ = 0
        config["psnr_"] = psnr_
        config["falseRate"] = falseRate_

        url_src = url_for('static', filename=config.get("originalPath"))
        url_dst = url_for('static', filename=config.get("stegPath"))
        
        config.update({'flag':'3', 'url_src':url_src, 'url_dst':url_dst, 'info_result':str_info, 'beforeLong':beforeLong,'beforeWide':beforeWide, 'beforeChannel':beforeChannel, 'afterLong':afterLong, 'afterWide':afterWide,'afterChannel':afterChannel,'imgbefore':img[np.where(condition)].tolist(),'ids':ids,'imgafter':stegData[np.where(condition)].tolist()})
        with open(os.path.join(basepath, 'static', config["dataDir"], 'config.json'), 'w') as f:
            json.dump(config, f)
        return config

@app.route('/handle-lsb-hide', methods=['POST'])
@app.route('/handle-pvd-hide', methods=['POST'])
@app.route('/handle-dct-hide', methods=['POST'])
def hide_html():
    try:
        logging.basicConfig(filename='log.log', level=logging.INFO)
        assert request.method == 'POST'
        done, data = get_secrets_and_coverimg(request)
        if done:
            f, secrets_text = data
        else:
            return jsonify({"error": '1'}.update(data))
        basepath = os.path.dirname(__file__)
        cover, config = prepareFile(f)
        uri = str(request.url_rule)
        
        if uri == '/handle-lsb-hide':
            config["algName"] = "lsb"
            if config.get("fileType") == "wav":
                wav_path = os.path.join(basepath, "static", config.get("originalPath"))
                wav_frames = wav_lsb_hide(wav_path, secrets_text)
                
                def performance_test():
                    wav_frames = wav_lsb_hide(wav_path)
                exe_time = timeit.timeit(performance_test, number=100)
                config["exe_time"] = exe_time
                
                config = afterHide(wav_frames, config)
                
                url_src = url_for('static', filename=config.get("originalPath"))
                url_dst = url_for('static', filename=config.get("stegPath"))
                return jsonify({'error': '0', 'flag': '4', 'id': config["id"]})
            else:
                bits = request.form.get("bits")
                
                img_lsb,img = lsb_hide(cover, secrets_text, int(bits))
                
                def performance_test():
                    img_lsb,img = lsb_hide(cover, secrets_text, int(bits))
                # exe_time = timeit.timeit(performance_test, number=100)
                # config["exe_time"] = exe_time
                
                config = afterHide(img_lsb, config, img)
                return jsonify({'error': '0', 'flag': '4', 'id': config["id"]})
        elif uri == '/handle-pvd-hide':
            config["algName"] = "pvd"

            coverImgPath = os.path.join(basepath, "static", config["originalPath"])
            kuangdu = request.form.get("pvdkuangdu")
            message = " PVD隐写 "
            
            #TODO: 配置range width
            img_pvd = PVD().hide(coverImgPath, secrets_text)
            img = cover
            config = afterHide(img_pvd, config, img)
            return jsonify({'error': '0', 'flag': '4', 'id': config["id"]})
        elif uri == '/handle-dct-hide':
            config["algName"] = "dct"
            juzhen_kuang = request.form.get("juzhen_kuang")
            img_dct,img = dct_hide(cover, mes=secrets_text, filename="", N=int(juzhen_kuang))
            config = afterHide(img_dct, config, img)
            return jsonify({'error': '0', 'flag': '4', 'id': config["id"]})
    except Exception as e:
        print("隐藏过程中发生错误: ", str(e))
        logging.error(f"隐藏过程中发生错误: {str(e)}")
        return jsonify({"error": '1', 'message':"隐藏过程中发生错误，请检查文件和参数"})
    
@app.route('/handle-lsb-extract', methods=['POST'])
@app.route('/handle-pvd-extract', methods=['POST'])
@app.route('/handle-dct-extract', methods=['POST'])
def extract_html():
    try:
        assert request.method == 'POST'
        if request.files:
            f = request.files['file2']
            logging.info(f"接收文件: {f.filename}")
            print(f)
            if not (f and allowed_file(f.filename)):
                print("file error")
                logging.error("文件格式错误")
                # return render_template('lsb.html', message="上传文件格式错误")
                return jsonify({'flag':'1', 'error': '1', 'message': "文件格式错误"})
        else:
            logging.warning("没有文件上传")
            # return render_template('lsb.html', message="文件为空")
            return jsonify({'flag':'2', 'error': '1', 'message': "没有文件上传"})
        basepath = os.path.dirname(__file__)
        stego, config = prepareFile(f)
        uri = str(request.url_rule)
        
        if uri == '/handle-lsb-extract':
            config["algName"] = "lsb"
            if config.get("fileType") == "wav":
                wav_path = os.path.join(basepath, "static", config.get("originalPath"))
                secrets = wav_lsb_extract(wav_path)
                # return render_template("lsb.html", secrets=secrets)
                return jsonify({'flag':'3', 'secrets':secrets})
            else:
                bits = request.form.get("bits")
                secrets = lsb_reveal(stego, int(bits), "./static/output.txt")
                if secrets == 'error_of_bits!':
                    return jsonify({'flag':'2', 'error': '1', 'message': "没有文件上传"})
                # return render_template("lsb.html", secrets=secrets)
                return jsonify({'flag':'3', 'error':'0', 'secrets':secrets})
        elif uri == '/handle-pvd-extract':
            config["algName"] = "pvd"
            kuangdu = request.form.get("pvdkuangdu")
            #TODO: 配置range width
            stegoImgPath = Path(os.path.join(basepath, "static", config["originalPath"]))
            secrets = PVD().extract(stegoImgPath)
            if secrets == 'error_of_bits!':
                return jsonify({'flag':'2', 'error': '1', 'message': "pvd范围宽度选取不对"})
            return jsonify({'flag':'3', 'error':'0', 'secrets':secrets})
        elif uri == '/handle-dct-extract':
            config["algName"] = "dct"
            juzhen_kuang = request.form.get("juzhen_kuang")
            coefficient = request.form.get("dct_coefficient")
            secrets = dct_reveal(stego, "./static/output.txt")
            if secrets == 'error_of_bits!':
                return jsonify({'flag':'2', 'error': '1', 'message': "矩阵宽高选取不对"})
            return jsonify({'flag':'3', 'error':'0', 'secrets':secrets})
    except Exception as e:
        print("提取过程中发生问题:", str(e))
        logging.error(f"提取过程中发生错误: {str(e)}")
        return jsonify({'flag':'2', 'error': '1', 'message': "提取过程中发生错误，请检查文件和参数"})

@app.route('/lsb.html', methods=['GET', 'POST'])
def lsb():
    if request.method == 'GET':
        return render_template('lsb.html')

@app.route('/dct.html', methods=['GET', 'POST'])
def dct():
    if request.method == 'GET':
        return render_template('dct.html')

@app.route('/pvd.html', methods=['GET', 'POST'])
def pvd():
    if request.method == 'GET':
        return render_template('pvd.html')

@app.route('/progress')	#前端javascript靠这个路由来访问数据
def progress():

    @stream_with_context
    def generate():
        global ratio
        ratio = progress_bar_ratio

        while ratio < 200:			#循环获取，直到达到100%
            if ratio >= 1:  # 仅发送大于等于1的进度值
               yield "data:" + str(ratio) + "\n\n"
           
            ratio = progress_bar_ratio
            #设置睡眠时间，避免后端在处理一帧的过程中前端就访问了好多次
            time.sleep(0.2)			

    return Response(generate(), mimetype='text/event-stream')	#用数据流的方式发送给后端

@app.route('/image_ok.html', methods=['GET'])
def image_ok_html():
    configId = request.args.get('id')
    basepath = os.path.dirname(__file__)
    if configId is None:
        configFilePath = os.path.join(basepath, 'static/data/config.json')
    else:
        configFilePath = os.path.join(basepath, 'static/data', time.strftime("%Y.%m.%d"), configId, 'config.json')
    with open(configFilePath, 'r') as f:
        config = json.load(f)
        #TODO: Showing all pixels is to expensive for browser. Just pick the first 100 pixels now. Show the changed pixels may be better.
        config["ids"] = config["ids"][:100]
        config["imgafter"] = config["imgafter"][:100]
        config["imgbefore"] = config["imgbefore"][:100]
    return render_template('image_ok.html', show_information=True, **config)

@app.route('/error.html', methods=['GET'])
def error_html():
    return render_template('error.html')

@app.route('/extReading.html', methods=['GET'])
def extReading():
    if request.method == 'GET':
        return render_template('extReading.html')



@app.route('/analysis.html', methods=['GET', 'POST'])
def analysis():
    if request.method == 'GET':
        return render_template('analysis.html')
    else:
        f = request.files['file1']
        basepath = os.path.dirname(__file__)
        path = os.path.join(basepath, 'static/data/analysis', f.filename)
        f.save(path)
        
        # rs分析计算隐写率
        img = image.imread(path)
        data = np.array(img[:,:,0], dtype=np.int16)
        rate,rs0,rs1 = stegRate(data,4,4)
        rate_str = f'估计隐写率为：{rate*100}%'
        rs0_str = f'RM={rs0[0]}, SM={rs0[1]}, RM_={rs0[2]}, SM_={rs0[3]}'
        rs1_str = f'RM={rs1[0]}, SM={rs1[1]}, RM_={rs1[2]}, SM_={rs1[3]} (最低位翻转后)'


        # 卡方分析
        img = Image.open(path)
        width = img.size[0]
        height = img.size[1]
        img = img.convert("L")
        martix = np.array(img)
        p = stgPrb(martix[0:int(width), 0:int(height)])
        p_str =  f'卡方值：{p}'
        return render_template('analysis.html',message=rate_str,rs0=rs0_str,rs1=rs1_str,p=p_str)

@app.route('/', methods=['GET'])
@app.route('/main.html', methods=['GET'])
def main():
    if request.method == 'GET':
        return render_template('main.html')

@app.route('/navigation.html', methods=['GET', 'POST'])
def navigation():
    if request.method == 'GET':
        return render_template('navigation.html')


if __name__ == '__main__':
    # app.run(debug=True)
    app.run(host="0.0.0.0", port="5000")
